from django.core.management.base import BaseCommand, CommandError

from buckets.models import Bucket


class Command(BaseCommand):
    """
    设置存储读写权限
    """
    error_buckets = []
    help = """
            ** manage.py bucket_lock --lock=xxx [--all] [--id-gt=xxx] [--id-lt=xxx] **
            ** --lock=1 锁定写， 2 锁定读写， 0 正常读写 **
           """

    def add_arguments(self, parser):
        parser.add_argument(
            '--lock', default=None, dest='lock',
            help='bucket lock flag',
        )
        parser.add_argument(
            '--all', default=None, nargs='?', dest='all', const=True,  # 当命令行有此参数时取值const, 否则取值default
            help='Update column "lock" for all buckets',
        )

        parser.add_argument(
            '--id-gt', default=0, dest='id-gt', type=int,
            help='id > xxx.',
        )

        parser.add_argument(
            '--id-lt', default=0, dest='id-lt', type=int,
            help='id < xxx.',
        )

    def handle(self, *args, **options):
        print(options)
        buckets = self.get_buckets(**options)
        if buckets is None:
            return

        lock = int(options['lock'])

        if input('Are you sure you want to do this? \n\n'
                 + "Type 'yes' to continue, or 'no' to cancel: ") != 'yes':
            raise CommandError("cancelled.")

        self.modify_buckets(buckets, lock)

    def _get_buckets(self, **options):
        """
        获取给定的bucket或所有bucket
        :param options:
        :return:
        """
        all_ = options['all']
        id_gt = options['id-gt']
        id_lt = options['id-lt']

        qs = Bucket.objects.all()
        if qs is None:
            self.stdout.write(self.style.ERROR(
                f"bucket is None."))
            return None

        # 全部的桶
        if all_ is not None:
            self.stdout.write(self.style.NOTICE('All buckets.'))
            return qs
        elif id_gt > 0 and id_lt > 0:
            self.stdout.write(self.style.NOTICE(f'buckets:({id_gt}, {id_lt})'))
            qs = qs.filter(id__gt=id_gt, id__lt=id_lt)
        elif id_gt > 0:
            self.stdout.write(self.style.NOTICE(f'buckets:({id_gt}, end_id)'))
            qs = qs.filter(id__gt=id_gt)
        elif id_lt > 0:
            self.stdout.write(self.style.NOTICE(f'buckets:(start_id, {id_lt})'))
            qs = qs.filter(id__lt=id_lt)
        else:
            self.stdout.write(self.style.ERROR(
                f"[--all] [--id-gt=xxx] [--id-lt=xxx] is None."))
            return None

        return qs

    def get_buckets(self, **options):

        buckets = self._get_buckets(**options)
        if buckets is None:
            return None

        buckets = buckets.order_by('id')
        return buckets

    def modify_one_bucket(self, bucket, lock_id):
        """
        为一个bucket执行sql

        :param bucket: Bucket obj
        :return:
        pool_id 存在跳过  存在检测
        """
        ret = True

        try:
            bucket.set_lock(lock=lock_id)
        except Exception as e:
            ret = False
            self.stdout.write(self.style.ERROR(
                f"Bucket(id={bucket.id}, name={bucket.name}), error when set lock: {str(e)}"))

        if ret:
            self.stdout.write(self.style.SUCCESS(
                f"Bucket(id={bucket.id}, name={bucket.name})，success to set lock; bucket lock: {bucket.lock}"))

        return ret

    def modify_buckets(self, buckets, lock_id):
        """
        单线程为bucket执行sql
        :param buckets:
        :return: None
        """
        count = 0
        for bucket in buckets:
            ok = self.modify_one_bucket(bucket, lock_id=lock_id)
            if ok:
                count += 1
            else:
                self.stdout.write(self.style.ERROR(
                    f"Bucket(id={bucket.id}, name={bucket.name}), error when set lock."))
                break

        self.stdout.write(
            self.style.SUCCESS(
                'Successfully exec sql for {0} buckets'.format(count)
            )
        )
